package org.lanyu.springainovel.chat.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.servlet.http.HttpServletResponse;
import lombok.AllArgsConstructor;
import org.lanyu.springainovel.common.config.ChatModelManager;
import org.lanyu.springainovel.common.entity.AiModel;
import org.lanyu.springainovel.chat.vo.ImagePromptVO;
import org.lanyu.springainovel.common.vo.RestVO;
import org.lanyu.springainovel.common.service.AiModelService;
import org.springframework.ai.image.*;
import org.springframework.ai.openai.OpenAiImageOptions;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;

import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;

@Tag(name = "图片生成", description = "文生图，图生图")
@RestController
@RequestMapping("/image")
@AllArgsConstructor
public class ImageChatController {

    private final ChatModelManager chatModelManager;
    private final AiModelService aiModelService;

    @GetMapping("/v1")
    @Operation(summary = "单图生成", description = "直接返回Response")
    public void image(HttpServletResponse response, ImagePromptVO imagePromptVO) {
        // 通过模型ID查找模型配置
        AiModel modelConfig = aiModelService.findById(imagePromptVO.getModelId());
        if (modelConfig == null) {
            throw new RuntimeException("未找到指定模型ID: " + imagePromptVO.getModelId());
        }
        
        ImageOptions options = OpenAiImageOptions.builder()
                .model(modelConfig.getName())
                .N(imagePromptVO.getNum() != null ? imagePromptVO.getNum() : 1)
                .width(imagePromptVO.getWidth() != null ? imagePromptVO.getWidth() : 512)
                .height(imagePromptVO.getHeight() != null ? imagePromptVO.getHeight() : 512)
                .build();
        ImageMessage imageMessage = new ImageMessage(imagePromptVO.getMessage(), imagePromptVO.getWeight());
        List<ImageMessage> messageList = new ArrayList<>();
        messageList.add(imageMessage);
        ImagePrompt imagePrompt = new ImagePrompt(messageList, options);

        // 通过ChatModelManager获取生图模型
        ImageModel imageModel = chatModelManager.getOrCreateImageModel(modelConfig);
        ImageResponse imageResponse = imageModel.call(imagePrompt);
        String imageUrl = imageResponse.getResult().getOutput().getUrl();

        try {
            URL url = URI.create(imageUrl).toURL();
            InputStream in = url.openStream();

            response.setHeader("Content-Type", MediaType.IMAGE_PNG_VALUE);
            response.getOutputStream().write(in.readAllBytes());
            response.getOutputStream().flush();
        } catch (IOException e) {
            response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
        }
    }

    @PostMapping("/v2")
    @Operation(summary = "多图生成", description = "返回所有图片的url")
    public RestVO<List<String>> image2(@RequestBody ImagePromptVO imagePromptVO) {
        // 通过模型ID查找模型配置
        AiModel modelConfig = aiModelService.findById(imagePromptVO.getModelId());
        if (modelConfig == null) {
            throw new RuntimeException("未找到指定模型ID: " + imagePromptVO.getModelId());
        }
        
        ImageOptions options = OpenAiImageOptions.builder()
                .model(modelConfig.getName())
                .N(imagePromptVO.getNum() != null ? imagePromptVO.getNum() : 1)
                .width(imagePromptVO.getWidth() != null ? imagePromptVO.getWidth() : 512)
                .height(imagePromptVO.getHeight() != null ? imagePromptVO.getHeight() : 512)
                .build();
        ImageMessage imageMessage = new ImageMessage(imagePromptVO.getMessage(), imagePromptVO.getWeight());
        List<ImageMessage> messageList = new ArrayList<>();
        messageList.add(imageMessage);
        ImagePrompt imagePrompt = new ImagePrompt(messageList, options);
        
        // 通过ChatModelManager获取生图模型
        ImageModel imageModel = chatModelManager.getOrCreateImageModel(modelConfig);
        ImageResponse imageResponse = imageModel.call(imagePrompt);
        
        List<ImageGeneration> imageGenerations = imageResponse.getResults();
        List<String> imageUrls = new ArrayList<>();
        for (ImageGeneration imageGeneration : imageGenerations) {
            String a = imageGeneration.getOutput().getUrl();
            imageUrls.add(a);
        }
        return RestVO.success(imageUrls);
    }
}
